import pandas as pd
import numpy as np
from math import cos, sin, pi
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d


def matrix_rotation(accel, orient):
    accel_shape = accel.shape
    orient_shape = orient.shape

    if accel_shape != orient_shape:
        print('shape does not match!')
        pass

    row = accel_shape[0]
    res = np.zeros([row, 3])
    for i in range(0, row):
        x, y, z = accel.X[i], accel.Y[i], accel.Z[i]
        conv_const = pi / 180
        z_theta, x_theta, y_theta = orient.X[i] * conv_const, -orient.Y[i] * conv_const, -orient.Z[i] * conv_const
        x_rot = np.array([
            [1, 0, 0, 0],
            [0, cos(x_theta), -sin(x_theta), 0],
            [0, sin(x_theta), cos(x_theta), 0],
            [0, 0, 0, 1]
        ])
        y_rot = np.array([
            [cos(y_theta), 0, -sin(y_theta), 0],
            [0, 1, 0, 0],
            [sin(y_theta), 0, cos(y_theta), 0],
            [0, 0, 0, 1]
        ])
        z_rot = np.array([
            [cos(z_theta), -sin(z_theta), 0, 0],
            [sin(z_theta), cos(z_theta), 0, 0],
            [0, 0, 1, 0],
            [0, 0, 0, 1]
        ])
        theta = np.array([x_theta, y_theta, z_theta, 1]).T
        prev = np.array([x, y, z, 1]).T
        curr = np.dot(x_rot, prev)
        curr = np.dot(y_rot, curr)
        curr = np.dot(z_rot, curr)
        res[i, :] = curr[0: 3]
    return res


def visualize_sensorlog(tot):
    fig = plt.figure()
    ax = plt.axes(projection='3d')
    color_map = ['Greens', 'Blues', 'Reds', 'Oranges', 'Purples']

    for i in range(0, tot):
        accel = pd.read_csv('sensorlog_accel_' + str(i+1) + '.csv')
        orient = pd.read_csv('sensorlog_orient_' + str(i+1) + '.csv')

        res = matrix_rotation(accel, orient)

        xdata, ydata, zdata = res[:, 0], res[:, 1], res[:, 2]
        ax.scatter3D(xdata, ydata, zdata, c=zdata, cmap=color_map[i])

    plt.show()


visualize_sensorlog(5)

# accel = pd.read_csv('sensorlog_accel_1.csv')
# orient = pd.read_csv('sensorlog_orient_1.csv')
#
# res = matrix_rotation(accel, orient)
#
# print(res)